import type { BasicTableProps } from "../types/table";
import { ROW_KEY } from "../const";
import { useMemo, useState } from "react";

export function useTableExpand(propsRef: BasicTableProps, tableData: Recordable[]) {
	const [expandedRowKeys, setExpandedRowKeys] = useState<(string | number)[]>([]);

	const getAutoCreateKey = useMemo(() => propsRef.autoCreateKey && !propsRef.rowKey, [propsRef.autoCreateKey, propsRef.rowKey]);

	const getRowKey = useMemo(() => {
		return getAutoCreateKey ? ROW_KEY : propsRef.rowKey;
	}, [propsRef.rowKey, getAutoCreateKey]);

	const getExpandOption = useMemo(() => {
		if (!propsRef.isTreeTable) return {};
		return {
			expandedRowKeys,
			onExpandedRowsChange: (keys: string[]) => {
				setExpandedRowKeys(keys);
			}
		};
	}, [propsRef.isTreeTable, expandedRowKeys]);

	function expandAll() {
		setExpandedRowKeys(getAllKeys());
	}

	function expandRows(keys: (string | number)[]) {
		// use row ID expands the specified table row
		const { isTreeTable } = propsRef;
		if (!isTreeTable) return;
		setExpandedRowKeys([...expandedRowKeys, ...keys]);
	}

	function getAllKeys(data?: Recordable[]) {
		const keys: string[] = [];
		const { childrenColumnName } = propsRef;
		(data || tableData).forEach(item => {
			keys.push(item[getRowKey as string]);
			const children = item[childrenColumnName || "children"];
			if (children?.length) keys.push(...getAllKeys(children));
		});
		return keys;
	}

	function collapseAll() {
		setExpandedRowKeys([]);
	}

	return { getExpandOption, expandAll, expandRows, collapseAll };
}
